import pandas as pd

# 读取数据

# 定义文件路径
moneyflow_path = r'F:\Personal\data\merged\moneyflow.csv'
limit_list_path = r'F:\Personal\data\merged\limit_list_d.csv'
limit_list_df = pd.read_csv(limit_list_path)
moneyflow_df = pd.read_csv(moneyflow_path)

# 确保日期列是 datetime 类型
limit_list_df['trade_date'] = pd.to_datetime(limit_list_df['trade_date'])
moneyflow_df['trade_date'] = pd.to_datetime(moneyflow_df['trade_date'])

# 初始化结果列表
filtered_moneyflow_list = []

# 遍历 limit_list_df 中的每一行
for _, limit_row in limit_list_df.iterrows():
    ts_code = limit_row['ts_code']
    trade_date = limit_row['trade_date']

    # 找到 moneyflow_df 中对应 ts_code 且日期在 trade_date 前后 60 天的数据
    start_date = trade_date - pd.Timedelta(days=60)
    end_date = trade_date + pd.Timedelta(days=60)
    filtered_df = moneyflow_df[(moneyflow_df['ts_code'] == ts_code) &
                               (moneyflow_df['trade_date'] >= start_date) &
                               (moneyflow_df['trade_date'] <= end_date)].copy()

    # 添加 limit_list_df 中的所有列
    for col in limit_list_df.columns:
        filtered_df[col] = limit_row[col]

    # 将过滤后的数据添加到结果列表
    filtered_moneyflow_list.append(filtered_df)

# 将结果列表合并为一个 DataFrame
filtered_moneyflow_df = pd.concat(filtered_moneyflow_list, ignore_index=True)

# 保存结果到一个新的 CSV 文件
filtered_moneyflow_df.to_csv('filtered_moneyflow_with_limit_info.csv', index=False)

print("Filtered moneyflow data saved to 'filtered_moneyflow_with_limit_info.csv'")